package com.wan.admin;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.lang.StringUtils;
import org.apache.flume.Context;
import org.apache.flume.Event;
import org.apache.flume.EventDeliveryException;
import org.apache.flume.PollableSource;
import org.apache.flume.conf.Configurable;
import org.apache.flume.event.SimpleEvent;
import org.apache.flume.instrumentation.SourceCounter;
import org.apache.flume.source.AbstractPollableSource;
import org.apache.rocketmq.client.consumer.DefaultMQPullConsumer;
import org.apache.rocketmq.client.consumer.PullResult;
import org.apache.rocketmq.client.consumer.PullStatus;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.text.SimpleDateFormat;
import java.util.*;


public class RocketMQSource extends AbstractPollableSource implements Configurable {

	private static final Logger LOG = LoggerFactory.getLogger(RocketMQSource.class);

	/**
	 * RocketMQ主题(消息类型)
	 */
	private String topic;

	/**
	 * RocketMQ标签(消息类型中的二级分类)
	 */
	private String tags;

	/**
	 * RocketMQ主题头部名称配置项
	 */
	private String topicHeaderName;

	/**
	 * RocketMQ标签头部名称配置项
	 */
	private String tagsHeaderName;

	/**
	 * 一次最多拉取条数
	 */
	private Integer batchSize;

	/**
	 * 自定义参数
	 */
	private String p1;
	private String p2;

	/**
	 * 消息拉取消费者
	 */
	private DefaultMQPullConsumer consumer;

	/**
	 * 统计
	 */
	private SourceCounter sourceCounter;

	/**
	 * 是否输出自定义日志
	 */
	private Boolean isPrintLog;



	@Override
	public void doConfigure(Context context) {
		isPrintLog = context.getBoolean(RocketMQSourceUtil.IS_PRINT_LOG, true);

		if (isPrintLog){
			System.out.println("开始进行配置启动RocketMQ-Flume-Source");
		}
		topic = Preconditions.checkNotNull(context.getString(RocketMQSourceUtil.TOPIC_CONFIG,RocketMQSourceUtil.TOPIC_DEFAULT), "必须指定RocketMQ分片, 例如: a1.sources.r1.topic=TestTopic");
		tags = context.getString(RocketMQSourceUtil.TAGS_CONFIG, RocketMQSourceUtil.TAGS_DEFAULT);
		topicHeaderName = context.getString(RocketMQSourceUtil.TOPIC_HEADER_NAME_CONFIG, RocketMQSourceUtil.TOPIC_HEADER_NAME_DEFAULT);
		tagsHeaderName = context.getString(RocketMQSourceUtil.TAGS_HEADER_NAME_CONFIG, RocketMQSourceUtil.TAGS_HEADER_NAME_DEFAULT);
		batchSize = context.getInteger(RocketMQSourceUtil.BATCH_SIZE_CONFIG, RocketMQSourceUtil.BATCH_SIZE_DEFAULT);
		p1 = context.getString(RocketMQSourceUtil.P1_CONFIG, RocketMQSourceUtil.P1_DEFAULT);
		p2 = context.getString(RocketMQSourceUtil.P2_CONFIG, RocketMQSourceUtil.P2_DEFAULT);
		if (isPrintLog){
			System.out.println("加载参数成功,topic:"+topic+";tags:"+tags+";batchSize:"+batchSize+";p1:"+p1+";p2:"+p2);
		}
		consumer = Preconditions.checkNotNull(RocketMQSourceUtil.getConsumer(context));
		//注册监听队列
		consumer.registerMessageQueueListener(topic, null);
		if (isPrintLog){
			System.out.println("消费者成功注册:"+ JSON.toJSONString(consumer));
		}

		//设置Source统计
		if (sourceCounter == null){
			sourceCounter = new SourceCounter(getName());
			if (isPrintLog){
				System.out.println("统计对象创建成功:"+ JSON.toJSONString(sourceCounter));
			}
		}
	}


	/**
	 * 开始拉取消息
	 */
	@Override
	public synchronized void doStart() {
		if (isPrintLog){
			System.out.println("开始启动消费者!");
		}

		try {
			// 启动Consumer
			consumer.start();
		} catch (MQClientException e) {
			if (isPrintLog){
				System.out.println("启动消费者失败!");
			}
			LOG.error("开始拉取RocketMQ资源失败", e);
			Throwables.propagate(e);
		}

		if (isPrintLog){
			System.out.println("启动统计数据源对象!");
		}
		sourceCounter.start();
	}


	/**
	 * 读取数据处理
	 * @return 读取状态
	 * @throws EventDeliveryException 事件输送异常
	 */
	@Override
	public PollableSource.Status doProcess() throws EventDeliveryException {
		//新建事件列表
		List<Event> eventList = Lists.newArrayList();
		Map<MessageQueue, Long> offsetMap = Maps.newHashMap();
		Event event;
		Map<String, String> headers;

		if (isPrintLog){
			System.out.println("读取数据处理程序开始!");
		}

		try
		{
			//从指定topic中拉取所有消息队列
			Set<MessageQueue> queues = Preconditions.checkNotNull(consumer.fetchSubscribeMessageQueues(topic));
			if (isPrintLog){
				System.out.println("从指定主题中拉取所有的消息队列对象:"+JSON.toJSONString(queues));
			}
			for (MessageQueue queue : queues) {
				//获取offset进度
				long offset = getMessageQueueOffset( queue );
				if (isPrintLog){
					System.out.println("迭代队列————>消费者offset进度:"+offset);
				}

				//消费者拉取结果
				PullResult pullResult = consumer.pullBlockIfNotFound(queue,tags,offset,batchSize);
				//进度集合存入(无论状态是否是FOUND都记录)
				offsetMap.put(queue,pullResult.getNextBeginOffset());

				if (isPrintLog){
					System.out.println("迭代队列————>消费者拉取结果:"+JSON.toJSONString(pullResult));
				}

				if (pullResult.getPullStatus() == PullStatus.FOUND) {
					if (isPrintLog){
						System.out.println("迭代队列————>拉取状态为FOUND");
					}
					//取出拉取结果消息列表
					for (MessageExt messageExt : pullResult.getMsgFoundList()) {
						//新建简单事件
						event = new SimpleEvent();
						headers = new HashMap<String,String>();
						//将消息的主题和标签存入头部
						headers.put( topicHeaderName,messageExt.getTopic() );
						headers.put( tagsHeaderName,messageExt.getTags() );

						//获取事件内容文本
						String body = new String( messageExt.getBody(),"UTF-8" );
						if (LOG.isDebugEnabled()) {
							LOG.debug( "MessageQueue={}, Topic={}, Tags={}, Message: {}",new Object[] {
									queue,messageExt.getTopic(),messageExt.getTags(),body} );
						}
						if (isPrintLog){
							System.out.println("迭代队列————>在迭代的消息数据————>MessageQueue:"+JSON.toJSONString(queue)+";Topic:"+JSON.toJSONString(messageExt.getTopic())
									+";Tags:"+JSON.toJSONString(messageExt.getTags())+";Message:"+JSON.toJSONString(body));
						}

						//设置赋值自定义配置
						String appId = null;
						String createTime = null;
						try {
							JSONObject jb = JSONObject.parseObject( body );
							SimpleDateFormat format = new SimpleDateFormat( "yyyyMMdd" );
							appId = jb.getString( p1 );
							createTime = format.format( jb.getLong( p2 ) );
						}
						catch (Exception e) {
							LOG.error( "JSON解析异常,Message:{}",e.getMessage(),e );
							if (isPrintLog){
								System.out.println("迭代队列————>JSON解析异常————>Message:"+JSON.toJSONString(e));
							}
						}

						if (StringUtils.isBlank( appId )) {
							appId = "unknown";
						}
						if (StringUtils.isBlank( createTime )) {
							SimpleDateFormat format = new SimpleDateFormat( "yyyyMMdd" );
							createTime = format.format( new Date() );
						}
						headers.put( p1,appId );
						headers.put( p2,createTime );
						//将消息内容存入事件内容
						event.setBody( messageExt.getBody() );
						event.setHeaders( headers );

						if (isPrintLog){
							System.out.println("迭代队列————>将消息存入事件中:"+JSON.toJSONString(event));
						}

						eventList.add( event );
						if (isPrintLog){
							System.out.println("迭代队列————>事件列表:"+JSON.toJSONString(eventList));
						}
					}
					if (isPrintLog){
						System.out.println("迭代队列————>存储进度:"+JSON.toJSONString(offsetMap));
					}
				}
			}

			//如果事件列表不为空
			if (eventList.size()>0) {
				if (isPrintLog){
					System.out.println("事件列表不为空处理!");
				}
				sourceCounter.incrementAppendBatchReceivedCount();
				sourceCounter.addToEventReceivedCount(eventList.size());
				getChannelProcessor().processEventBatch(eventList);
				sourceCounter.incrementAppendBatchAcceptedCount();
				sourceCounter.addToEventAcceptedCount(eventList.size());
				eventList.clear();
				if (isPrintLog){
					System.out.println("事件列表不为空处理结束:"+JSON.toJSONString(eventList));
				}
			}

			for (Map.Entry<MessageQueue, Long> entry : offsetMap.entrySet()) {
				//更新offset
				this.putMessageQueueOffset(entry.getKey(), entry.getValue());
			}

		}
		catch (Exception e)
		{
			LOG.error("RocketMQSource消费者消息异常", e);
			if (isPrintLog){
				System.out.println("RocketMQSource消费者消息异常:"+JSON.toJSONString(e));
			}
			return Status.BACKOFF;
		}
		return Status.READY;
	}


	/**
	 * 停止拉取消息
	 */
	@Override
	public synchronized void doStop() {
		//统计停止
		sourceCounter.stop();
		//停止Consumer
		consumer.shutdown();
	}


	/**
	 * 获取消息队列消费进度
	 * @param queue 消息队列
	 * @return 进度
	 * @throws MQClientException MQ客户端异常
	 */
	private long getMessageQueueOffset(MessageQueue queue) throws MQClientException {
		// 从Broker获取Offset
		long offset = consumer.fetchConsumeOffset(queue, false);
		if (offset < 0L) {
			offset = 0L;
		}

		if (isPrintLog){
			System.out.println("获取消息队列消费进度:"+offset);
		}
		return offset;
	}


	/**
	 * 存储拉取进度
	 * @param queue 消息队列
	 * @param offset 进度
	 * @throws MQClientException 消息
	 */
	private void putMessageQueueOffset(MessageQueue queue, long offset) throws MQClientException {
		//存储Offset,客户端每隔5s会定时刷新到Broker或者写入本地缓存文件
		consumer.updateConsumeOffset(queue, offset);
		if (isPrintLog){
			System.out.println("存储拉取进度:"+offset);
		}
	}

}
